#include "kernel_operator.h"
constexpr int32_t BUFFER_NUM = 1; // tensor num for each queue
constexpr int32_t DATA_BLOCK = 32;
template<typename TYPE_X> class KernelRenorm;

// float类型的直接renorm操作
template <>
class KernelRenorm<float>
{
public:
    __aicore__ inline KernelRenorm() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR z, float p, float max_norm, uint32_t dim, RenormCustomTilingData tiling)
    {
      	this->tiling = tiling;
      	this->p = p;
        this->max_norm = max_norm;
        this->dim = dim;
        this->block_idx = AscendC::GetBlockIdx();
		this->tileNum = tiling.tileNum;
        this->tileLength = tiling.tileLength / BUFFER_NUM;
        this->blockLength = tiling.blockLength;
        this->totalLength = tiling.totalLength;
        this->rawNum = tiling.rawNum;
        this->columnNum = tiling.columnNum;
        this->EvenCore = tiling.EvenCore;
		this->tailNum = tiling.tailNum;
        this->alignFlag = tiling.alignFlag;

        if(this->dim == 0){
            if(this->EvenCore == this->block_idx && this->alignFlag == 1){
        		xGm.SetGlobalBuffer((__gm__ float *)x + this->blockLength * this->block_idx, this->blockLength + this->tileLength * (this->tailNum - this->tileNum));
        		zGm.SetGlobalBuffer((__gm__ float *)z + this->blockLength * this->block_idx, this->blockLength + this->tileLength * (this->tailNum - this->tileNum));
          		this->tileNum = this->tailNum;
          	}
          	else if (this->EvenCore != this->block_idx && this->alignFlag == 1){
        		xGm.SetGlobalBuffer((__gm__ float *)x + this->blockLength * this->block_idx, this->blockLength);
        		zGm.SetGlobalBuffer((__gm__ float *)z + this->blockLength * this->block_idx, this->blockLength);
          	}
            else{
                xGm.SetGlobalBuffer((__gm__ float *)x + this->block_idx * (uint64_t)this->blockLength);
        	    zGm.SetGlobalBuffer((__gm__ float *)z + this->block_idx * (uint64_t)this->blockLength);
                if(this->EvenCore == this->block_idx){
            	    this->tileNum = this->tailNum;
                }
          }
        }
        else if(this->dim == 1){
          	this->tileLength = this->tileLength * DATA_BLOCK / sizeof(float);
            xGm.SetGlobalBuffer((__gm__ float *)x + this->block_idx * (uint64_t)this->tileNum);
        	zGm.SetGlobalBuffer((__gm__ float *)z + this->block_idx * (uint64_t)this->tileNum);
            if(this->EvenCore == this->block_idx){
              	this->tileNum = this->tailNum;
            }
        }

        sum_params.inner = this->tileLength;
        sum_params.outter = 1;
        sum_params.n = this->tileLength;
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(float));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmpBuffer_power, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmpBuffer_sum, sum_params.outter * sizeof(float));
        pipe.InitBuffer(tmpBuffer_norm, sum_params.outter * sizeof(float));

    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++)
        {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
      	AscendC::LocalTensor<float> xLocal = inQueueX.AllocTensor<float>();
        if(this->dim == 0){
            if (this->alignFlag == 1){
                  AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
            }
            else{
                  AscendC::DataCopyExtParams copyParams{(uint16_t)(1), (uint32_t)(this->columnNum * sizeof(float)),
                                                0, 0, 0};
                  AscendC::DataCopyPadExtParams<float> padParams{true, 0, (uint8_t)(this->tileLength - this->columnNum), 0};
        	      AscendC::DataCopyPad(xLocal, this->xGm[progress * this->columnNum], copyParams, padParams);
            }
      	}
        else if(this->dim == 1){
			AscendC::DataCopyExtParams copyParams{(uint16_t)this->rawNum, 1 * sizeof(float),
                                                (uint16_t)((this->columnNum - 1) * sizeof(float)), 0, 0};
        	AscendC::DataCopyPadExtParams<float> padParams{true, 0, DATA_BLOCK/sizeof(float) - 1, 0};
        	AscendC::DataCopyPad(xLocal, this->xGm[progress], copyParams, padParams);
        }
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
      	AscendC::LocalTensor<float> xLocal = inQueueX.DeQue<float>();
        AscendC::LocalTensor<float> zLocal = outQueueZ.AllocTensor<float>();
        AscendC::LocalTensor<float> tmp_power = tmpBuffer_power.Get<float>();
        AscendC::LocalTensor<float> tmp_sum = tmpBuffer_sum.Get<float>();
        AscendC::LocalTensor<float> tmp_norm = tmpBuffer_norm.Get<float>();

		int16_t index = 0;
        // AscendC::PRINTF("x0: %f", xLocal.GetValue(index));
        // AscendC::DumpTensor(xLocal, 0, this->tileLength);

		// AscendC::DumpTensor(tmp, this->block_idx, this->tileLength );
        AscendC::Power<float, false>(tmp_power, xLocal, this->p);
        // AscendC::PRINTF("tmp_power: %f", tmp_power.GetValue(index));
		// AscendC::DumpTensor(tmp_power, this->block_idx, this->tileLength);
        AscendC::Sum(tmp_sum, tmp_power, this->sum_params);
        // AscendC::PRINTF("tmp_sum: %f", tmp_sum.GetValue(index));
		// AscendC::DumpTensor(tmp_sum, this->block_idx, this->tileLength);
        AscendC::Power(tmp_norm, tmp_sum, (float)1.0 / this->p, this->sum_params.outter);
        float norm = tmp_norm.GetValue(index);
        AscendC::PRINTF("norm: %f", norm);
        if (norm < this->max_norm)
        {
            AscendC::Muls(zLocal, xLocal, (float)1.0, this->tileLength);
        }
        else
        {
            AscendC::Muls(zLocal, xLocal, (float)this->max_norm / norm, this->tileLength);
        }

		// AscendC::DumpTensor(tmp_output, this->block_idx, this->tileLength );
		// AscendC::DumpTensor(zLocal, this->block_idx, this->tileLength );



        outQueueZ.EnQue<float>(zLocal);
        inQueueX.FreeTensor(xLocal);

    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
      	AscendC::LocalTensor<float> zLocal = outQueueZ.DeQue<float>();
        if(this->dim == 0){
            if (this->alignFlag == 1){
                  AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
            }
            else{
                  AscendC::DataCopyExtParams copyParams{(uint16_t)1, (uint32_t)(this->columnNum * sizeof(float)),
                                                0, 0, 0};
        	      AscendC::DataCopyPad(this->zGm[progress * this->columnNum], zLocal, copyParams);
//                  AscendC::DataCopy(zGm[progress * this->columnNum], zLocal, this->tileLength);
            }
      	}
        else if(this->dim == 1){
			AscendC::DataCopyExtParams copyParams{(uint16_t)this->rawNum, 1 * sizeof(float), 0, (uint16_t)((this->columnNum - 1) * sizeof(float)), 0};
        	AscendC::DataCopyPad(this->zGm[progress], zLocal, copyParams); // 从VECIN->GM搬运40Bytes
        }
        outQueueZ.FreeTensor(zLocal);

    }

private:

    AscendC::GlobalTensor<float> xGm;
    AscendC::GlobalTensor<float> zGm;

    RenormCustomTilingData tiling;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lastTileLength;
    uint32_t totalLength;
    uint32_t rawNum;
    uint32_t columnNum;
    uint32_t EvenCore;
	uint32_t tailNum;
    uint32_t alignFlag;

    float p;
    float max_norm;
    uint32_t dim;
    int32_t block_idx;
    int32_t pad_till_length;

    AscendC::SumParams sum_params;

    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer_power;
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer_sum;
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer_norm;
};

// 通用模板：支持half\bfloat16类型的直接renorm操作
template <typename TYPE_X>
class KernelRenorm
{
public:
    __aicore__ inline KernelRenorm() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR z, float p, float max_norm, uint32_t dim, RenormCustomTilingData tiling)
    {
      	this->tiling = tiling;
      	this->p = p;
        this->max_norm = max_norm;
        this->dim = dim;
        this->block_idx = AscendC::GetBlockIdx();
		this->tileNum = tiling.tileNum;
        this->tileLength = tiling.tileLength / BUFFER_NUM;
        this->blockLength = tiling.blockLength;
        this->totalLength = tiling.totalLength;
        this->rawNum = tiling.rawNum;
        this->columnNum = tiling.columnNum;
        this->EvenCore = tiling.EvenCore;
		this->tailNum = tiling.tailNum;
        this->alignFlag = tiling.alignFlag;

        if(this->dim == 0){
          if(this->EvenCore == this->block_idx && this->alignFlag == 1){
        		xGm.SetGlobalBuffer((__gm__ TYPE_X *)x + this->blockLength * this->block_idx, this->blockLength + this->columnNum * (this->tailNum - this->tileNum));
        		zGm.SetGlobalBuffer((__gm__ TYPE_X *)z + this->blockLength * this->block_idx, this->blockLength + this->columnNum * (this->tailNum - this->tileNum));
          		this->tileNum = this->tailNum;
          }
          else if (this->EvenCore != this->block_idx && this->alignFlag == 1){
        		xGm.SetGlobalBuffer((__gm__ TYPE_X *)x + this->blockLength * this->block_idx, this->blockLength);
        		zGm.SetGlobalBuffer((__gm__ TYPE_X *)z + this->blockLength * this->block_idx, this->blockLength);
          }
          else{
                xGm.SetGlobalBuffer((__gm__ TYPE_X *)x + this->block_idx * (uint64_t)this->blockLength);
        	    zGm.SetGlobalBuffer((__gm__ TYPE_X *)z + this->block_idx * (uint64_t)this->blockLength);
                if(this->EvenCore == this->block_idx){
            	    this->tileNum = this->tailNum;
                }
          }

        }
        else if(this->dim == 1){
          	this->tileLength = this->tileLength * DATA_BLOCK / sizeof(TYPE_X);
            xGm.SetGlobalBuffer((__gm__ TYPE_X *)x + this->block_idx * (uint64_t)this->tileNum);
        	zGm.SetGlobalBuffer((__gm__ TYPE_X *)z + this->block_idx * (uint64_t)this->tileNum);
            if(this->EvenCore == this->block_idx){
            	this->tileNum = this->tailNum;
            }
        }

        sum_params.inner = this->tileLength;
        sum_params.outter = 1;
        sum_params.n = this->tileLength;
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(TYPE_X));
        pipe.InitBuffer(outQueueZ, BUFFER_NUM, this->tileLength * sizeof(TYPE_X));
        pipe.InitBuffer(tmpBuffer, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmpBuffer_power, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmpBuffer_sum, sum_params.outter * sizeof(float));
        pipe.InitBuffer(tmpBuffer_norm, sum_params.outter * sizeof(float));
        pipe.InitBuffer(tmpBuffer_output, this->tileLength * sizeof(float));

    }
    __aicore__ inline void Process()
    {
        int32_t loopCount = this->tileNum * BUFFER_NUM;
        for (int32_t i = 0; i < loopCount; i++)
        {
            CopyIn(i);
            Compute(i);
            CopyOut(i);
        }
    }

private:
    __aicore__ inline void CopyIn(int32_t progress)
    {
      	AscendC::LocalTensor<TYPE_X> xLocal = inQueueX.AllocTensor<TYPE_X>();
        if(this->dim == 0){
            if (this->alignFlag == 1){
                  AscendC::DataCopy(xLocal, xGm[progress * this->tileLength], this->tileLength);
            }
            else{
                  AscendC::DataCopyExtParams copyParams{(uint16_t)(1), (uint32_t)(this->columnNum * sizeof(TYPE_X)),
                                                0, 0, 0};
                  AscendC::DataCopyPadExtParams<TYPE_X> padParams{true, 0, (uint8_t)(this->tileLength - this->columnNum), 0};
        	      AscendC::DataCopyPad(xLocal, this->xGm[progress * this->columnNum], copyParams, padParams);
            }

      	}
        else if(this->dim == 1){
			AscendC::DataCopyExtParams copyParams{(uint16_t)this->rawNum, 1 * sizeof(TYPE_X),
                                                (uint16_t)((this->columnNum - 1) * sizeof(TYPE_X)), 0, 0};
        	AscendC::DataCopyPadExtParams<TYPE_X> padParams{true, 0, DATA_BLOCK/sizeof(TYPE_X) - 1, 0};
        	AscendC::DataCopyPad(xLocal, this->xGm[progress], copyParams, padParams);
        }
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t progress)
    {
      	AscendC::LocalTensor<TYPE_X> xLocal = inQueueX.DeQue<TYPE_X>();
        AscendC::LocalTensor<TYPE_X> zLocal = outQueueZ.AllocTensor<TYPE_X>();
        AscendC::LocalTensor<float> tmp = tmpBuffer.Get<float>();
        AscendC::LocalTensor<float> tmp_power = tmpBuffer_power.Get<float>();
        AscendC::LocalTensor<float> tmp_sum = tmpBuffer_sum.Get<float>();
        AscendC::LocalTensor<float> tmp_norm = tmpBuffer_norm.Get<float>();
        AscendC::LocalTensor<float> tmp_output = tmpBuffer_output.Get<float>();

		int16_t index = 0;
//         AscendC::PRINTF("column: %d", this->columnNum);
//         AscendC::PRINTF("tileLength: %d", this->tileLength);
//         AscendC::PRINTF("x0: %f", xLocal.GetValue(index));
//         AscendC::DumpTensor(xLocal, this->block_idx, this->tileLength);
    	AscendC::Cast(tmp, xLocal, AscendC::RoundMode::CAST_NONE, this->tileLength);

//		 AscendC::DumpTensor(tmp, this->block_idx, this->tileLength );
        AscendC::Power<float, false>(tmp_power, tmp, this->p);
//         AscendC::PRINTF("tmp_power: %f", tmp_power.GetValue(index));
//		 AscendC::DumpTensor(tmp_power, this->block_idx, this->tileLength);
        AscendC::Sum(tmp_sum, tmp_power, this->sum_params);
//         AscendC::PRINTF("tmp_sum: %f", tmp_sum.GetValue(index));
//		 AscendC::DumpTensor(tmp_sum, this->block_idx, this->tileLength);
        AscendC::Power(tmp_norm, tmp_sum, (float)1.0 / this->p, this->sum_params.outter);
        float norm = tmp_norm.GetValue(index);
        AscendC::PRINTF("norm: %f", norm);
        if (norm < this->max_norm)
        {
            AscendC::Muls(tmp_output, tmp, (float)1.0, this->tileLength);
        }
        else
        {
            AscendC::Muls(tmp_output, tmp, (float)this->max_norm / norm, this->tileLength);
        }

        AscendC::Cast(zLocal, tmp_output, AscendC::RoundMode::CAST_ODD, this->tileLength);

		// AscendC::DumpTensor(tmp_output, this->block_idx, this->tileLength );
//		 AscendC::DumpTensor(zLocal, this->block_idx, this->tileLength );



        outQueueZ.EnQue<TYPE_X>(zLocal);
        inQueueX.FreeTensor(xLocal);

    }

    __aicore__ inline void CopyOut(int32_t progress)
    {
      	AscendC::LocalTensor<TYPE_X> zLocal = outQueueZ.DeQue<TYPE_X>();
        if(this->dim == 0){
            if (this->alignFlag == 1){
                  AscendC::DataCopy(zGm[progress * this->tileLength], zLocal, this->tileLength);
            }
            else{
                  AscendC::DataCopyExtParams copyParams{(uint16_t)1, (uint32_t)(this->columnNum * sizeof(TYPE_X)),
                                                0, 0, 0};
        	      AscendC::DataCopyPad(this->zGm[progress * this->columnNum], zLocal, copyParams);
//                  AscendC::DataCopy(zGm[progress * this->columnNum], zLocal, this->tileLength);
            }

      	}
        else if(this->dim == 1){
			AscendC::DataCopyExtParams copyParams{(uint16_t)this->rawNum, 1 * sizeof(TYPE_X), 0, (uint16_t)((this->columnNum - 1) * sizeof(TYPE_X)), 0};
        	AscendC::DataCopyPad(this->zGm[progress], zLocal, copyParams);
        }
        outQueueZ.FreeTensor(zLocal);

    }

private:

    AscendC::GlobalTensor<TYPE_X> xGm;
    AscendC::GlobalTensor<TYPE_X> zGm;

    RenormCustomTilingData tiling;
    uint32_t blockLength;
    uint32_t tileNum;
    uint32_t tileLength;
    uint32_t lastTileLength;
    uint32_t totalLength;
    uint32_t rawNum;
    uint32_t columnNum;
    uint32_t EvenCore;
	uint32_t tailNum;
	uint32_t alignFlag;

    float p;
    float max_norm;
    uint32_t dim;
    int32_t block_idx;
    int32_t pad_till_length;

    AscendC::SumParams sum_params;

    AscendC::TPipe pipe;
    AscendC::TQue<AscendC::TPosition::VECIN, BUFFER_NUM> inQueueX;
    AscendC::TQue<AscendC::TPosition::VECOUT, BUFFER_NUM> outQueueZ;
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer;
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer_power;
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer_sum;
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer_norm;
    AscendC::TBuf<AscendC::TPosition::VECCALC> tmpBuffer_output;
};




extern "C" __global__ __aicore__ void renorm_custom(GM_ADDR x, GM_ADDR z, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
	uint32_t dim = static_cast<uint32_t>(tiling_data.dim);
	float p = static_cast<float>(tiling_data.p);
	float max_norm = static_cast<float>(tiling_data.max_norm);
//	reinterpret_cast<RenormCustomTilingData *>(tiling)

    if constexpr (std::is_same_v<DTYPE_X, bfloat16_t>)
    {
        KernelRenorm<DTYPE_X> op;
        op.Init(x, z, p, max_norm, dim, tiling_data);
        op.Process();
    }
    else if constexpr (std::is_same_v<DTYPE_X, float16_t>)
    {
        KernelRenorm<DTYPE_X> op;
        op.Init(x, z, p, max_norm, dim, tiling_data);
        op.Process();
    }
    else if constexpr (std::is_same_v<DTYPE_X, float>)
    {
        KernelRenorm<float> op;
        op.Init(x, z, p, max_norm, dim, tiling_data);
        op.Process();
    }
    else
    {
        return;
    }
}